import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as data
import torchvision
import matplotlib.pyplot as plt

EPOCH = 1
BATCH_SIZE = 50
DATA_DIR = '/data1/alpha/ML/mnist2'
LR = 0.001

DOWNLOAD_MNIST = True

train_data = torchvision.datasets.MNIST(
    root=DATA_DIR,
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST
)

# print(train_data.train_data.size())
# print(train_data.train_labels.size())
#
# plt.imshow(train_data.train_data[0].numpy(), cmap='gray')
# plt.title('%i' % train_data.train_labels[0])
# plt.show()


train_loader = data.DataLoader(dataset=train_data,
                               batch_size=BATCH_SIZE,
                               shuffle=True,
                               num_workers=2)

test_data = torchvision.datasets.MNIST(
    root=DATA_DIR,
    train=False,
)

test_x = Variable(torch.unsqueeze(test_data.test_data, dim=1), volatile=True).type(torch.FloatTensor)[0: 2000]
test_y = test_data.test_labels[0:2000]


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # (1, 28, 28)
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1,
                      out_channels=16,
                      kernel_size=5,
                      stride=1,
                      padding=2,  # if  stride =1, padding =(kernel-1)/2
                      ), # (16, 28, 28)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2) # (16, 14, 14)
        )

        # (16, 14, 14)
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16,
                      out_channels=32,
                      kernel_size=5,
                      stride=1,
                      padding=2,  # if  stride =1, padding =(kernel-1)/2
                      ),  # (32, 14, 14)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2) # (32, 7, 7)
        )

        self.linear = nn.Linear(32 * 7 * 7, 10)

        self.log_softmax = nn.LogSoftmax()

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)   # batch, 32, 7, 7
        x = x.view(x.size(0), -1)  # batch, 32 * 7 * 7
        x = self.linear(x)
        output = self.log_softmax(x)
        return output


cnn = CNN()

opt = torch.optim.Adam(cnn.parameters(), lr=LR)
loss_func = nn.CrossEntropyLoss()

for epoch in range(EPOCH):
    print(epoch)
    for step, (batch_x, batch_y) in enumerate(train_loader):
        b_x = Variable(batch_x)
        b_y = Variable(batch_y)

        output = cnn(b_x)
        loss = loss_func(output, b_y)
        opt.zero_grad()
        loss.backward()
        opt.step()

        if step % 50 == 0:
            test_output = cnn(test_x)
            pred_y = torch.max(test_output, 1)[1].data.squeeze()
            accuracy = sum(pred_y == test_y) / test_y.size(0)
            print('Epoch: ', epoch,
                  '| train_loss: %.4f' % loss.data[0],
                  '| test_accuracy: ' , accuracy)
